
import math


def calculate_ndcg_for_sample(predicted_list, ground_truth_item, k=10):

    # 截取前k个预测结果
    predicted_list = predicted_list[:k]
    # 计算相关性分数列表
    relevance = []
    for item_id in predicted_list:
        if item_id == ground_truth_item:
            relevance.append(1)  # 相关项
        else:
            relevance.append(0)  # 不相关项
    # 计算DCG@k
    dcg = 0
    for i, rel in enumerate(relevance):
        # 位置i的折损因⼦为log2(i+2)
        discount = math.log2(i + 2)
        dcg += rel / discount
    # 计算IDCG@k
    # 在本任务中，理想情况是将唯⼀相关项放在第⼀位
    idcg = 1 / math.log2(1 + 1)  # = 1
    # 计算NDCG@k
    if idcg > 0:
        ndcg = dcg / idcg
    else:
        ndcg = 0
    return ndcg


# 使⽤⽰例
# 假设模型预测的电影ID排序为[111, 1893, 684, 2492, ...]
# ⽤⼾实际观看的下⼀部电影ID为1893
if __name__ == '__main__':
    predicted_list = [111, 1893, 684, 2492, 3654, 2422, 176, 1629, 229, 3155]
    ground_truth_item = 1893
    ndcg10 = calculate_ndcg_for_sample(predicted_list, ground_truth_item, k=10)
    print(f"NDCG@10 = {ndcg10}")
    # 输出: NDCG@10 = 0.63093
